package com.github.kezhenxu94.data.sink;

import org.apache.kafka.clients.producer.KafkaProducer;
import org.apache.kafka.clients.producer.ProducerRecord;

import java.util.Properties;
import java.util.concurrent.*;
import java.util.function.Function;

/**
 * Created by kezhenxu on 1/5/17.
 *
 * @author kezhenxu (kezhenxu94@163.com)
 */
public class KafkaSink<K, V> implements Sink<V> {

    protected Properties properties;
    protected KafkaProducer<K, V> kafkaProducer;
    protected BlockingQueue<V> queue;
    protected String topic;
    protected ExecutorService executor;
    protected Function<V, K> keySupplier;

    private volatile boolean isRunning;

    @Override
    public void setup() throws Exception {
        executor = Executors.newSingleThreadExecutor();
        properties = new Properties();
        properties.load(getClass().getResourceAsStream("/kafka-sink.properties"));
        kafkaProducer = new KafkaProducer<>(properties);
        queue = new LinkedBlockingQueue<>(200);
    }

    @Override
    public void start() throws Exception {
        isRunning = true;
        executor.execute(() -> {
            while (isRunning()) {
                try {
                    V toPut = queue.take();
                    getLogger().debug("Putting {} to kafka. ", toPut);
                    kafkaProducer.send(new ProducerRecord<>(topic, keySupplier.apply(toPut), toPut));
                } catch (InterruptedException e) {
                    getLogger().error(e.getMessage(), e);
                }
            }
        });
    }

    @Override
    public void close() throws Exception {
        try {
            isRunning = false;
            while (isRunning()) {
                TimeUnit.MILLISECONDS.sleep(500L);
            }
            executor.awaitTermination(500L, TimeUnit.MILLISECONDS);
        } finally {
            if (kafkaProducer != null) {
                kafkaProducer.close();
                kafkaProducer = null;
            }
        }
    }

    @Override
    public boolean isRunning() {
        return isRunning || !queue.isEmpty();
    }

    @Override
    public void put(V v) throws Exception {
        if (!isRunning()) {
            throw new IllegalStateException("Sink has been closed. ");
        }
        queue.offer(v);
    }

    public String getTopic() {
        return topic;
    }

    public void setTopic(String topic) {
        this.topic = topic;
    }

    public Function<V, K> getKeySupplier() {
        return keySupplier;
    }

    public void setKeySupplier(Function<V, K> keySupplier) {
        this.keySupplier = keySupplier;
    }
}
